"""
Memory management callbacks for continual learning.
"""

from typing import Optional, List

from ..task_callbacks import TaskCallback


class MemoryManagementCallback(TaskCallback):
    """
    Callback for managing memory resources during continual learning.

    Args:
        clear_features (bool): Whether to clear cached features
        clear_gradients (bool): Whether to clear gradients after updates
        memory_efficient_training (bool): Whether to use memory-efficient training
        max_cached_tasks (int): Maximum number of previous tasks to keep in cache
    """

    def __init__(
        self,
        clear_features: bool = True,
        clear_gradients: bool = True,
        memory_efficient_training: bool = False,
        max_cached_tasks: int = 2,
    ):
        super().__init__()
        self.clear_features = clear_features
        self.clear_gradients = clear_gradients
        self.memory_efficient_training = memory_efficient_training
        self.max_cached_tasks = max_cached_tasks
        self.cached_tasks = []

    def on_task_end(self, trainer, pl_module, task_id, dataloader=None):
        """Clean up resources when a task ends."""
        super().on_task_end(trainer, pl_module, task_id, dataloader)

        # Clear feature cache if not needed for next task
        if self.clear_features and hasattr(pl_module.model, "backbone"):
            backbone = pl_module.model.backbone
            if hasattr(backbone, "clear_cache"):
                backbone.clear_cache()

        # Manage task caches
        self.cached_tasks.append(task_id)
        if len(self.cached_tasks) > self.max_cached_tasks:
            # Remove oldest task from cache
            oldest_task = self.cached_tasks.pop(0)
            self._clear_task_specific_resources(pl_module, oldest_task)

    def on_after_backward(self, trainer, pl_module):
        """Clean up after backward pass."""
        if self.clear_gradients and not pl_module.automatic_optimization:
            # For models where we want to manually control gradient retention
            for name, param in pl_module.named_parameters():
                if not name.startswith("model.adapters") and param.grad is not None:
                    param.grad = None

    def _clear_task_specific_resources(self, pl_module, task_id):
        """Clear resources specific to a task."""
        # Clear any components that cache task-specific data
        if hasattr(pl_module, "components"):
            for component in pl_module.components:
                if hasattr(component, "clear_task_cache"):
                    component.clear_task_cache(task_id)

        # Clear task-specific metrics
        if hasattr(pl_module, "test_acc") and str(task_id) in pl_module.test_acc:
            del pl_module.test_acc[str(task_id)]
